import torch
from torch.utils.data import Dataset, DataLoader
from tokenizer.gpt2 import GPT2Tokenizer

class GPT2DatasetV1(Dataset):
    def __init__(self, text, tokenizer, max_length, stride):
        self.tokenizer = tokenizer
        self.input_ids = []
        self.targets = []

        tokens = tokenizer.encode(text)
        for i in range(0, len(tokens) - max_length, stride):
            self.input_ids.append(tokens[i:i+max_length])
            self.targets.append(tokens[i+1:i+max_length+1])

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return torch.tensor(self.input_ids[idx]), torch.tensor(self.targets[idx])

def create_gpt2_dataloader_v1(
        text,
        max_length=256,
        stride=128,
        batch_size=4,
        shuffle=True,
        drop_last=True,
        num_workers=0):
    tokenizer = GPT2Tokenizer()
    dataset = GPT2DatasetV1(text, tokenizer, max_length, stride)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)
    return dataloader
